import cv2
import subprocess
import argparse
import threading
import queue
import signal
from func_timeout import func_set_timeout, FunctionTimedOut

from logger import get_logger
from alert import alert_thread_wrapper
from detect import *


parser = argparse.ArgumentParser()
# 流地址
parser.add_argument('--input', type=str, default='rtsp://admin:a123456789@10.42.0.94:554/Streaming/Channels/101?transportmode=unicast&profile=Profile_1', help="输入视频流（OpenCV参数）")
parser.add_argument('--output', type=str, default='rtmp://localhost/live/camera-24:32:ae:c0:52:04', help="输出视频流（FFMpeg参数）")
# 报警参数
parser.add_argument('--alert_collect_url', type=str, help="报警的url", default='http://localhost:5000/alert/collect/camera-24:32:ae:c0:52:04')
parser.add_argument('--alert_sensitivity', type=float, help='报警的灵敏度，数值越高越灵敏', default=0.3)
parser.add_argument('--enter_alert_objs', type=str, nargs='+', help='报警的类型，有物体进入时会触发报警', default=None)
parser.add_argument('--leave_alert_objs', type=str, nargs='+', help='报警的类型，有物体离开时会触发报警', default=None)
parser.add_argument('--linkage_alert_objs', type=str, nargs='+', help='报警的类型，当这些类型的目标同时出现在画面中时才会触发联动报警', default=None)
parser.add_argument('--pass_alert_objs', type=str, nargs='+', help='报警的类型，当这些类型的目标经过画面后才会报警，另外需要同时指定经过的方向才会报警（和经过方向配合使用）', default=None)
parser.add_argument('--pass_alert_directions', type=str, nargs='+', help='物体经过的方向，只有物体按照指定方向经过时才会报警（和经过报警的类型配合使用）', default=None)
parser.add_argument('--gather_alert_objs', type=str, nargs='+', help='报警的类型，每次该类型的物体在画面中的数量超过阈值时都会触发报警', default=None)
parser.add_argument('--gather_alert_threshhold', type=int, help='当画面中物体的数量达到（大于等于）该阈值时才会触发报警', default=None)
# 区域入侵参数
parser.add_argument('--ai_region', type=float, nargs='+', action='append', help='报警的区域，按照(x1,y1,y2,y2...)的顺序传入N个点的坐标，该参数可以传入多次，以标注多区域', default=None)
parser.add_argument('--ai_region_highlight', action='store_true', help='是否将报警的区域高亮显示')
# 转播参数
parser.add_argument('--resolution', type=str, help="转播的分辨率(宽*高)，如640*480", default=None)
parser.add_argument('--cbr', type=int, help='转播的码率(单位为kbps)，如1024', default=1024)
parser.add_argument('--fps', type=int, help='转播的帧率，如20', default=None)
# 画面参数
parser.add_argument('--enable_draw_video_box', action='store_true', help="启用视频画面绘制识别框")
parser.add_argument('--disable_draw_alert_box', action='store_true', help="禁止报警图像绘制识别框")
# 线程数
parser.add_argument('--thread_num', default=3, type=int, help="线程数量，太多太少都会导致画面变卡以及延迟增高")

config, _ = parser.parse_known_args()

logger = get_logger('nvrpro.main')
logger.info(f'启动参数：{config}')

cap = cv2.VideoCapture(int(config.input) if config.input.isdigit() else config.input)
if not cap.isOpened():
    logger.fatal(f'无法读取设备的视频流[{config.input}]，请检查能否成功连通设备，以及设备的视频流能否正常播放')
    exit(1)
fps = config.fps if config.fps else cap.get(cv2.CAP_PROP_FPS)
if fps > 120:
    logger.fatal(f'获取帧率过高：{fps}，判断拉流出现了问题，无法正常拉流')
    exit(1)

# 是否改变分辨率
w, h = map(int, config.resolution.split('*')) if config.resolution else (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
resize = True if config.resolution else False

def process_arg_list(l):
    ret = []
    for i in l:
        if i is not None:
            ret.append(i) if type(i) == str else ret.extend(i)
    return ret

ffmpeg = subprocess.Popen(process_arg_list([
    'ffmpeg',
    '-y',
    '-fflags', 'nobuffer',
    '-f', 'rawvideo',
    '-vcodec','rawvideo',
    '-pix_fmt', 'bgr24',  # 我试过把ffmpeg的像素格式设为rgb24，结果就是帧率降低，可能BGR是ffmpeg的标准格式快一点
    '-s', f'{w}x{h}',
    '-r', str(fps),
    '-i', '-',
    '-c:v', 'libx264',
    ('-b:v', f"{config.cbr}k") if config.cbr else None,
    '-pix_fmt', 'yuv420p',
    '-preset', 'ultrafast',
    '-f', 'flv',
    config.output]), stdin=subprocess.PIPE)

alert_proc_queue = queue.Queue()
rknn_pool = RKNNPool(config.thread_num)

def exit_proc():
    ffmpeg.terminate()
    exit(0)

def pull_frame():
    global origin_frame, cap
    # 加个超时包装
    @func_set_timeout(5)
    def read():
        return cap.read()

    while cap.isOpened():
        try:
            ret, frame = read()  # 这个函数如果卡太久说明有问题，直接重连试试
        except FunctionTimedOut:
            ret, frame = False, None
            logger.info('读取视频流超时（5s内无法读取一帧）')
        if not ret:  # frame可能为None，这里加个判断
            logger.info('读取视频流出现问题，即将尝试重新读取视频流...')
            cap = cv2.VideoCapture(config.input)
            if not cap.isOpened() or cap.get(cv2.CAP_PROP_FPS) > 120:
                logger.fatal('拉流出现问题，无法正常获取设备rtsp视频流')
                exit_proc()  # 拉流还有问题就不重试了，直接退出，因为这个问题大概率不可能在几秒内解决，循环重试还浪费性能，就直接让nvr-manager过一段时间再重启程序吧
            else:
                logger.info('重新读取视频流成功')
            continue
        if frame is None:
            continue
        rknn_pool.detect_queue.put(frame)
    logger.info('摄像头关闭')
    exit_proc()

# 设置检测区域
if config.ai_region:
    logger.info(f'已启用区域检测，检测区域为{config.ai_region}')
    set_mask(config.ai_region, (w, h))
    if config.ai_region_highlight:
        logger.info(f'区域检测高亮已启用')

if config.enable_draw_video_box:
    logger.info('已启用绘制视频画面识别框')

if config.disable_draw_alert_box:
    logger.info('已禁用绘制报警图像识别框')

# 启动拉流线程，利用opencv不断获取视频的帧
threading.Thread(target=pull_frame, daemon=True).start()

# 启动报警检测线程，不断获取检测后的结果，并决定是否报警
threading.Thread(target=alert_thread_wrapper, args=(
    alert_proc_queue, 
    config.alert_collect_url, 
    config.disable_draw_alert_box,
    config.enter_alert_objs, 
    config.leave_alert_objs, 
    config.linkage_alert_objs, 
    config.pass_alert_objs,
    config.pass_alert_directions,
    config.gather_alert_objs,
    config.gather_alert_threshhold,
    config.alert_sensitivity
), daemon=True).start()

signal.signal(signal.SIGTERM, lambda _1, _2: exit_proc())

last_id = 0.
while True:
    id, image, results = rknn_pool.result_queue.get()
    if id < last_id:
        continue
    last_id = id
    alert_proc_queue.put((image, results))
    res_frame = image.copy()
    if config.ai_region_highlight:
        draw_mask(res_frame)
    if config.enable_draw_video_box:
        res_frame = draw_results(res_frame, results)
    if resize:
        res_frame = cv2.resize(res_frame, (w, h))
    ffmpeg.stdin.write(res_frame)
